/*
 * Copyright 2016-2017 TVI Go Easy.
 * Created on 2017/11/7 17:02
 */
package com.gitee.mechanic.core.utils;

import com.google.common.collect.Lists;
import com.trilead.ssh2.Connection;
import com.trilead.ssh2.SCPClient;
import com.trilead.ssh2.Session;
import com.gitee.mechanic.core.enums.HttpCode;
import com.gitee.mechanic.core.exception.SSHException;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.io.InputStream;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.List;

/**
 * SSH 工具类
 *
 * @author mechanic
 */
public class SSHUtils {

    private final Logger logger = LoggerFactory.getLogger(SSHUtils.class);
    /**
     * 主机名或IP
     */
    private String host;
    /**
     * 端口
     */
    private int port = 22;
    /**
     * 账户
     */
    private String account;
    /**
     * 密码
     */
    private String password;
    /**
     * 密钥
     */
    private String key;
    /**
     * 密钥的密码
     */
    private String keyPassword;
    /**
     * 连接超时时间(毫秒)
     */
    private int timeout = 3000;
    /**
     * 文件编码
     */
    private static final String ENCODING = "UTF-8";
    /**
     * 执行多个命令输出结果的分割线
     */
    private static final String SEPARATOR = "--SSH_EXECUTE_CMD_OUT_SPLIT-------------------------------------------";

    /**
     * SSH 连接
     */
    private Connection connection;

    public SSHUtils(String host, String account, String key, String keyPassword)  {
        this(host, 22, account, key, keyPassword);
    }

    public SSHUtils(String host, int port, String account, String key, String keyPassword)  {
        this(host, port, 3000, account, key, keyPassword);
    }

    public SSHUtils(String host, int port, int timeout, String account, String key, String keyPassword)
       {
        this.host = host;
        this.port = port;
        this.timeout = timeout;
        this.account = account;
        this.key = key;
        this.keyPassword = keyPassword;
        this.connection = new Connection(host, port);
    }

    public SSHUtils(String host, String account, String password)  {
        this(host, 22, account, password);
    }

    public SSHUtils(String host, int port, String account, String password)  {
        this(host, port, 3000, account, password);
    }

    public SSHUtils(String host, int port, int timeout, String account, String password)  {
        this.host = host;
        this.port = port;
        this.timeout = timeout;
        this.account = account;
        this.password = password;
        this.connection = new Connection(host, port);
    }

    public Connection getConnection() {
        return new Connection(host, port);
    }

    /**
     * 验证是否可登录服务器，会关闭Connection
     *
     * @return boolean
     */
    public boolean auth()  {
        boolean isAuth = false;
        try {
            isAuth = this.connect();
        } finally {
            if (connection != null) {
                connection.close();
            }
        }
        return isAuth;
    }

    /**
     * 执行命令，命令执行完毕后会关闭连接
     *
     * @param cmd 命令
     * @return 返回命令正常执行的屏幕输出信息
     */
    public List<List<String>> execute(String cmd)  {
        List<String> cmds = Lists.newArrayList();
        if (StringUtils.endsWith(cmd, ";")) {
            cmd = StringUtils.substringBeforeLast(cmd, ";");
        }
        cmds = Lists.newArrayList(StringUtils.split(cmd, ";"));
        return execute(cmds);
    }

    /**
     * 执行命令，命令执行完毕后会关闭连接
     *
     * @param cmds 命令集
     * @return 返回命令正常执行的屏幕输出信息
     */
    public List<List<String>> execute(String[] cmds)  {
        return execute(Lists.newArrayList(cmds));
    }

    /**
     * 执行命令，命令执行完毕后会关闭连接
     *
     * @param cmds 命令集
     * @return 返回命令正常执行的屏幕输出信息
     */
    public List<List<String>> execute(List<String> cmds)  {
        if (CollectionUtils.isEmpty(cmds)) {
            return Lists.newArrayList();
        }
        StringBuilder sb = new StringBuilder();
        String echo = "echo " + SEPARATOR + ";";
        for (String cmd : cmds) {
            cmd = StringUtils.trimToEmpty(cmd);
            if (StringUtils.isNotBlank(cmd)) {
                sb.append(echo);
                sb.append(cmd);
                if (!StringUtils.endsWith(cmd, ";")) {
                    sb.append(";");
                }
            }
        }

        //建立连接
        this.connect();
        Session session = null;
        try {
            //打开SSH会话
            session = connection.openSession();
            //执行命令
            session.execCommand(sb.toString());
            //读取命令执行后的错误信息

            String errorMessage = getMessageAsString(session.getStderr());
            //如果有错误信息则抛出异常
            if (StringUtils.isNotEmpty(errorMessage)) {
                throw getSSHException(sb, errorMessage);
            }

            return getMessageAsList(session.getStdout());

        } catch (IOException | SSHException e) {
            logger.error(e.toString());
            throw getSSHException(sb, "执行命令时发生错误");
        } finally {
            if (session != null) {
                session.close();
            }
            if (connection != null) {
                connection.close();
            }
        }
    }

    /**
     * 通过SSH将远程主机上的文件复制到本地
     *
     * @param remoteFiles          远程主机上的文件完整路径
     * @param localTargetDirectory 本地文件夹完整路径
     * @return String[]
     */
    public String[] pull(String[] remoteFiles, String localTargetDirectory)  {
        String[] localFiles = new String[remoteFiles.length];
        Path p = Paths.get(localTargetDirectory);
        SCPClient scpClient = new SCPClient(this.connection);
        this.connect();
        try {
            if (Files.exists(p)) {
                if (!Files.isDirectory(p)) {
                    Files.delete(p);
                }
            } else {
                Files.createDirectories(p);
            }
            scpClient.get(remoteFiles, localTargetDirectory);
            for (int i = 0; i < remoteFiles.length; i++) {
                Path fileName = Paths.get(remoteFiles[i]).getFileName();
                localFiles[i] = localTargetDirectory + "/" + fileName.toString();
            }
        } catch (IOException e) {
            logger.error(e.getCause().toString());
            throw getSSHException("复制文件失败");
        } finally {
            connection.close();
        }
        return localFiles;
    }

    /**
     * 通过SSH将本地文件批量上传到远程主机
     *
     * @param localFiles            本地文件完整路径
     * @param remoteTargetDirectory 远程主机上文件夹完整路径
     */
    public void push(String[] localFiles, String remoteTargetDirectory)  {
        SCPClient scp = new SCPClient(this.connection);
        this.connect();
        try {
            scp.put(localFiles, remoteTargetDirectory);
        } catch (IOException e) {
            logger.error(e.getCause().toString());
            throw getSSHException("上传文件失败");
        } finally {
            connection.close();
        }
    }

    /**
     * 建立SSH连接并通过带密码的密钥登录，不会关闭Connection
     *
     * @return boolean
     */
    private boolean connect()  {
        if (StringUtils.isNotBlank(this.key)) {
            return this.authWithKey();
        } else {
            return this.authWithPassword();
        }
    }

    /**
     * 建立SSH连接并通过带密码的密钥登录
     *
     * @return boolean
     */
    private boolean authWithKey()  {
        File file = Paths.get(key).toFile();
        try {
            this.connection.connect(null, this.timeout, 0);
            if (!this.connection.authenticateWithPublicKey(this.account, file, this.keyPassword)) {
                throw getSSHException("SSH Public Key 或 密码错误");
            }
        } catch (IOException e) {
            logger.error(e.toString());
            throw getSSHException("SSH 无法连接服务器");
        }
        return true;
    }

    /**
     * 建立SSH连接并通过密码登录
     *
     * @return boolean
     */
    private boolean authWithPassword()  {
        try {
            this.connection.connect(null, this.timeout, 0);
            if (!connection.authenticateWithPassword(this.account, this.password)) {
                throw getSSHException("SSH 账户或密码错误");
            }
        } catch (IOException e) {
            logger.error(e.toString());
            throw getSSHException("SSH 无法连接服务器");
        }
        return true;
    }

    /**
     * 将流形式的输出信息转为List
     *
     * @param stream 屏幕输出信息
     * @return List&lt;List&lt;String&gt;&gt;
     */
    private List<List<String>> getMessageAsList(InputStream stream)  {
        List<String> outList = null;
        try {
            outList = IOUtils.readLines(stream, ENCODING);
        } catch (IOException e) {
            logger.error(e.toString());
            throw getSSHException("读取 SSH 输出信息时发生错误");
        }

        List<List<String>> resultList = Lists.newLinkedList();
        if (CollectionUtils.isEmpty(outList)) {
            return resultList;
        }
        List<String> singleCmdOutList = Lists.newArrayList();
        for (int i = 0; i < outList.size(); i++) {
            String line = outList.get(i);
            if (StringUtils.isBlank(line)) {
                continue;
            }
            line = StringUtils.trim(line);
            if (StringUtils.equalsIgnoreCase(line, SEPARATOR)) {
                if (i > 0) {
                    resultList.add(singleCmdOutList);
                }
                singleCmdOutList = Lists.newArrayList();
            } else {
                singleCmdOutList.add(line);
            }
            if (i == (outList.size() - 1)) {
                resultList.add(singleCmdOutList);
            }
        }
        return resultList;
    }

    /**
     * 将流形式的输出信息转为String
     *
     * @param stream 屏幕输出信息
     * @return String
     */
    private String getMessageAsString(InputStream stream){
        StringBuilder messsage = new StringBuilder();
        try {
            for (String line : IOUtils.readLines(stream, ENCODING)) {
                messsage.append(line).append("\n");
            }
        } catch (IOException e) {
            logger.error(e.toString());
            throw getSSHException("读取 SSH 输出信息时发生错误");
        }
        return messsage.toString();
    }
    /**
     * 格式化异常信息
     *
     * @param error 异常信息
     * @return String
     */
    private SSHException getSSHException(CharSequence error) {
        return new SSHException(
          HttpCode.SSH_ERROR,
          SSHException.format(this.host, this.port, this.account, null, error)
        );
    }

    /**
     * 格式化带命令带异常信息
     *
     * @param cmd     命令
     * @param error 异常信息
     * @return String
     */
    private SSHException getSSHException(CharSequence cmd, CharSequence error) {
        return new SSHException(
          HttpCode.SSH_ERROR,
          SSHException.format(this.host, this.port, this.account, cmd, error)
        );
    }

    public String getHost() {
        return host;
    }

    public int getPort() {
        return port;
    }

    public String getAccount() {
        return account;
    }

    public String getPassword() {
        return password;
    }

    public String getKey() {
        return key;
    }

    public String getKeyPassword() {
        return keyPassword;
    }

    public int getTimeout() {
        return timeout;
    }
}
